#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EndoSight-UC AI推理服务器
使用真实的UCEIS评分模型进行内镜图像分析
"""

import os
import sys
import logging
import traceback
import time
from pathlib import Path
from typing import Dict, Any, Optional, List
import uuid

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from torchvision.models import ResNet50_Weights
from transformers import BertModel, BertTokenizer
from PIL import Image
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from pydantic import BaseModel

# 添加模型路径
MODEL_DIR = Path(r"D:\VSCodeProjects\UC_System\uc_model")
sys.path.insert(0, str(MODEL_DIR))

# 导入真实的模型
try:
    # 首先尝试导入兼容的多模态模型
    from inference_model_compat import EndoMultimodalModelCompat
    print("成功导入兼容的多模态EndoMultimodalModel")
    UCEISScoreModel = EndoMultimodalModelCompat
except ImportError as e:
    print(f"导入兼容模型失败: {e}")
    print("尝试导入原始多模态模型...")
    try:
        from models.fusion_model import EndoMultimodalModel
        print("成功导入多模态EndoMultimodalModel")
        UCEISScoreModel = EndoMultimodalModel
    except ImportError as e2:
        print(f"导入多模态模型失败: {e2}")
        print("尝试导入集成UCEISScoreModel...")
        try:
            from models.uceis_integrated_model import UCEISScoreModel, create_integrated_uceis_model
            print("成功导入集成UCEISScoreModel")
        except ImportError as e3:
            print(f"导入集成模型失败: {e3}")
            print("尝试导入原始UCEISScoreModel...")
            try:
                from models.uceis_model import UCEISScoreModel
                print("成功导入原始UCEISScoreModel")
            except ImportError as e4:
                print(f"导入原始模型也失败: {e4}")
                print("尝试创建基础模型...")
                try:
                    # 创建基础模型作为后备
                    import torch.nn as nn
                    class SimpleUCModel(nn.Module):
                        def __init__(self):
                            super().__init__()
                            self.backbone = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
                            num_ftrs = self.backbone.fc.in_features
                            self.backbone.fc = nn.Sequential(
                                nn.Linear(num_ftrs, 128),
                                nn.ReLU(),
                                nn.Dropout(0.5),
                                nn.Linear(128, 9)  # 0-8分
                            )

                        def forward(self, x):
                            logits = self.backbone(x)
                            return logits

                    UCEISScoreModel = SimpleUCModel
                    print("成功创建基础后备模型")
                except ImportError as e5:
                    print(f"创建后备模型也失败: {e5}")
                    sys.exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('ai_inference.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 全局变量
model = None
device = None
model_config = None

# 评分映射
SEVERITY_MAPPING = {
    0: "缓解期",
    1: "轻度",
    2: "轻度",
    3: "中度",
    4: "中度",
    5: "重度",
    6: "重度",
    7: "重度",
    8: "重度"
}

# 响应模型
class PredictionResponse(BaseModel):
    task_id: str
    uceis_score: int
    severity: str
    confidence: float
    detailed_scores: Dict[str, Any]
    processing_time: float

# 模型配置
def get_model_config() -> Dict[str, Any]:
    """获取模型配置"""
    return {
        'input_channels': 3,
        'feature_dims': [64, 128, 256, 512],
        'backbone_type': 'resnet50',
        'pretrained': True,
        'use_uncertainty': True,
        'vascular_weight': 1.0,
        'bleeding_weight': 1.0,
        'erosion_weight': 1.0,
        'total_weight': 1.0,
        'regression_weight': 0.5,
        'use_focal_loss': True,
        'focal_alpha': 1.0,
        'focal_gamma': 2.0
    }

def load_model() -> bool:
    """加载训练好的UCEIS模型"""
    global model, device, model_config

    try:
        logger.info("开始加载UCEIS模型...")

        # 检查GPU
        if torch.cuda.is_available():
            device = torch.device('cuda')
            logger.info(f"使用GPU: {torch.cuda.get_device_name(0)}")
        else:
            device = torch.device('cpu')
            logger.info("使用CPU进行推理")

        # 获取模型配置
        model_config = {
            'input_channels': 3,
            'feature_dims': [64, 128, 256, 512],
            'backbone_type': 'resnet50',
            'pretrained': True,
            'use_uncertainty': True,
            'vascular_weight': 1.0,
            'bleeding_weight': 1.0,
            'erosion_weight': 1.0,
            'total_weight': 1.0
        }

        # 创建模型实例
        logger.info("初始化模型结构...")

        # 使用兼容的模型结构，与我们训练的模型匹配
        logger.info("创建兼容的多模态模型结构...")
        class EndoMultimodalModelCompat(nn.Module):
            """兼容的模型结构，与我们训练的模型完全匹配"""
            def __init__(self):
                super().__init__()

                # 使用与训练时相同的模型结构
                logger.info("初始化兼容模型结构...")

                # ResNet-50特征提取器
                self.image_encoder = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
                self.image_encoder.fc = nn.Identity()

                # BERT文本编码器
                self.text_encoder = BertModel.from_pretrained('bert-base-uncased')

                # 特征投影层
                self.image_proj = nn.Linear(2048, 512)
                self.text_proj = nn.Linear(768, 512)

                # 融合层 - 完全匹配训练模型的结构
                self.fusion_layers = nn.ModuleDict({
                    '0': nn.Linear(512, 512),
                    '1': nn.BatchNorm1d(512),
                    '4': nn.Linear(512, 256),
                    '5': nn.BatchNorm1d(256),
                    '8': nn.Linear(256, 128),
                    '9': nn.BatchNorm1d(128),
                })

                # 最终输出层
                self.output_layer = nn.Sequential(
                    nn.Dropout(0.2),
                    nn.Linear(128, 64),
                    nn.Dropout(0.2),
                    nn.Linear(64, 8)  # 输出8个分数
                )

                logger.info("兼容模型结构初始化完成")

            def forward(self, image, text_inputs):
                # 图像特征提取
                image_features = self.image_encoder(image)
                image_features = self.image_proj(image_features)

                # 文本特征提取
                text_outputs = self.text_encoder(
                    input_ids=text_inputs['input_ids'],
                    attention_mask=text_inputs['attention_mask'],
                    return_dict=True
                )
                text_features = text_outputs.last_hidden_state[:, 0, :]
                text_features = self.text_proj(text_features)

                # 特征融合
                combined = image_features + text_features  # 相加融合，与训练时一致

                # 逐层通过融合层
                x = self.fusion_layers['0'](combined)
                x = self.fusion_layers['1'](x)
                x = torch.relu(x)
                x = self.fusion_layers['4'](x)
                x = self.fusion_layers['5'](x)
                x = torch.relu(x)
                x = self.fusion_layers['8'](x)
                x = self.fusion_layers['9'](x)
                x = torch.relu(x)

                # 最终输出
                output = self.output_layer(x)
                return output

        model = EndoMultimodalModelCompat()
        logger.info("兼容模型创建成功")

        model.to(device)

        # 查找最佳模型文件 - 优先使用新训练的模型
        model_paths = [
            Path(r"D:\VSCodeProjects\UC_System\uc_model\final_weights\uceis_best_model.pth"),  # 最终训练好的模型
        ]

        model_path = None
        for path in model_paths:
            if path.exists():
                model_path = path
                logger.info(f"找到模型文件: {model_path}")
                break

        # 如果找不到训练好的模型，使用预训练模型
        if model_path is None:
            logger.warning("未找到训练好的模型文件，使用预训练模型")
            # 不加载权重，直接使用预训练模型
            logger.info("使用预训练ResNet模型")
        else:
            # 加载模型权重
            logger.info(f"加载模型权重: {model_path}")
            try:
                checkpoint = torch.load(model_path, map_location=device)

                # 处理不同的checkpoint格式，允许部分匹配
                if isinstance(checkpoint, dict):
                    if 'model_state_dict' in checkpoint:
                        # 允许部分加载，只加载匹配的键
                        model_dict = model.state_dict()
                        checkpoint_dict = {k: v for k, v in checkpoint['model_state_dict'].items() if k in model_dict}
                        model_dict.update(checkpoint_dict)
                        model.load_state_dict(model_dict, strict=False)
                        logger.info("部分加载模型权重成功（忽略不匹配的层）")
                    elif 'state_dict' in checkpoint:
                        model_dict = model.state_dict()
                        checkpoint_dict = {k: v for k, v in checkpoint['state_dict'].items() if k in model_dict}
                        model_dict.update(checkpoint_dict)
                        model.load_state_dict(model_dict, strict=False)
                        logger.info("部分加载模型权重成功（忽略不匹配的层）")
                    else:
                        # 尝试直接加载权重，允许部分匹配
                        model_dict = model.state_dict()
                        checkpoint_dict = {k: v for k, v in checkpoint.items() if k in model_dict}
                        model_dict.update(checkpoint_dict)
                        model.load_state_dict(model_dict, strict=False)
                        logger.info("部分加载模型权重成功（忽略不匹配的层）")
                else:
                    # 直接加载，允许部分匹配
                    model_dict = model.state_dict()
                    checkpoint_dict = {k: v for k, v in checkpoint.items() if k in model_dict}
                    model_dict.update(checkpoint_dict)
                    model.load_state_dict(model_dict, strict=False)
                    logger.info("部分加载模型权重成功（忽略不匹配的层）")

                logger.info("成功加载训练好的模型权重")
            except Exception as e:
                logger.warning(f"加载权重失败，使用预训练模型: {e}")
                # 继续使用预训练模型

        # 设置为评估模式
        model.eval()

        logger.info("UCEIS模型加载成功！")
        logger.info(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")

        return True

    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        return False

def preprocess_image(image_path: str) -> torch.Tensor:
    """图像预处理"""
    try:
        # 加载图像
        image = Image.open(image_path).convert('RGB')
        logger.info(f"成功加载图像: {os.path.basename(image_path)}, 尺寸: {image.size}, 模式: {image.mode}")

        # 调整图像尺寸（保持宽高比）
        target_size = 512
        width, height = image.size
        aspect_ratio = width / height

        if width > height:
            new_width = target_size
            new_height = int(target_size / aspect_ratio)
        else:
            new_height = target_size
            new_width = int(target_size * aspect_ratio)

        image = image.resize((new_width, new_height), Image.LANCZOS)

        # 转换为tensor
        import torchvision.transforms as transforms
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ])

        image_tensor = transform(image).unsqueeze(0)  # 添加batch维度
        return image_tensor.to(device)

    except Exception as e:
        logger.error(f"图像预处理失败: {str(e)}")
        raise HTTPException(status_code=400, detail=f"图像预处理失败: {str(e)}")

def predict_with_model(image_tensor: torch.Tensor) -> Dict[str, Any]:
    """使用模型进行预测"""
    try:
        start_time = time.time()

        with torch.no_grad():
            # 根据模型类型进行预测
            if 'Multimodal' in model.__class__.__name__:
                # 多模态模型需要文本输入 - 使用默认文本
                from transformers import AutoTokenizer
                tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
                default_text = "Endoscopic image of colon"
                text_inputs = tokenizer(default_text, return_tensors='pt',
                                       padding=True, truncation=True, max_length=128)

                # 将文本输入移到设备
                text_inputs = {k: v.to(device) for k, v in text_inputs.items()}

                # 多模态预测
                outputs = model(images=image_tensor, text_inputs=text_inputs)
            else:
                # 单模态模型只需要图像
                outputs = model(image_tensor)

            # 处理不同的模型输出格式
            if isinstance(outputs, dict):
                # UCEISScoreModel输出格式
                if 'vascular_logits' in outputs and 'bleeding_logits' in outputs and 'erosion_logits' in outputs:
                    vascular_logits = outputs['vascular_logits']
                    bleeding_logits = outputs['bleeding_logits']
                    erosion_logits = outputs['erosion_logits']

                    # 获取预测的类别
                    vascular_score = torch.argmax(vascular_logits, dim=1).item()
                    bleeding_score = torch.argmax(bleeding_logits, dim=1).item()
                    erosion_score = torch.argmax(erosion_logits, dim=1).item()

                    # 计算总分
                    uceis_score = vascular_score + bleeding_score + erosion_score
                    uceis_score = max(0, min(8, int(uceis_score)))

                    # 获取置信度
                    vascular_confidence = torch.max(F.softmax(vascular_logits, dim=1), dim=1)[0].item()
                    bleeding_confidence = torch.max(F.softmax(bleeding_logits, dim=1), dim=1)[0].item()
                    erosion_confidence = torch.max(F.softmax(erosion_logits, dim=1), dim=1)[0].item()
                    confidence = (vascular_confidence + bleeding_confidence + erosion_confidence) / 3

                    # 详细评分
                    detailed_scores = {
                        'vascular_pattern': {
                            'score': int(vascular_score),
                            'confidence': float(vascular_confidence),
                            'description': ['正常', '纹理缺失', '严重缺失'][int(vascular_score)] if vascular_score <= 2 else '未知'
                        },
                        'bleeding': {
                            'score': int(bleeding_score),
                            'confidence': float(bleeding_confidence),
                            'description': ['无出血', '点状出血', '线状出血', '弥漫性出血'][int(bleeding_score)] if bleeding_score <= 3 else '未知'
                        },
                        'erosions_ulcers': {
                            'score': int(erosion_score),
                            'confidence': float(erosion_confidence),
                            'description': ['无糜烂', '糜烂<5mm', '糜烂>5mm', '溃疡形成'][int(erosion_score)] if erosion_score <= 3 else '未知'
                        }
                    }
                else:
                    # 其他字典格式
                    logits = list(outputs.values())[0]
                    uceis_score = torch.argmax(logits, dim=1).item()
                    uceis_score = max(0, min(8, int(uceis_score)))
                    confidence = torch.max(F.softmax(logits, dim=1), dim=1)[0].item()

                    # 生成模拟详细评分
                    detailed_scores = {
                        'vascular_pattern': {
                            'score': min(2, max(0, uceis_score // 3)),
                            'confidence': float(confidence),
                            'description': ['正常', '纹理缺失', '严重缺失'][min(2, max(0, uceis_score // 3))]
                        },
                        'bleeding': {
                            'score': min(3, max(0, uceis_score % 3)),
                            'confidence': float(confidence),
                            'description': ['无出血', '点状出血', '线状出血', '弥漫性出血'][min(3, max(0, uceis_score % 3))]
                        },
                        'erosions_ulcers': {
                            'score': min(3, max(0, uceis_score - (uceis_score // 3) - (uceis_score % 3))),
                            'confidence': float(confidence),
                            'description': ['无糜烂', '糜烂<5mm', '糜烂>5mm', '溃疡形成'][min(3, max(0, uceis_score - (uceis_score // 3) - (uceis_score % 3)))]
                        }
                    }
            else:
                # 直接输出logits
                logits = outputs
                uceis_score = torch.argmax(logits, dim=1).item()
                uceis_score = max(0, min(8, int(uceis_score)))
                confidence = torch.max(F.softmax(logits, dim=1), dim=1)[0].item()

                # 生成模拟详细评分
                detailed_scores = {
                    'vascular_pattern': {
                        'score': min(2, max(0, uceis_score // 3)),
                        'confidence': float(confidence),
                        'description': ['正常', '纹理缺失', '严重缺失'][min(2, max(0, uceis_score // 3))]
                    },
                    'bleeding': {
                        'score': min(3, max(0, uceis_score % 3)),
                        'confidence': float(confidence),
                        'description': ['无出血', '点状出血', '线状出血', '弥漫性出血'][min(3, max(0, uceis_score % 3))]
                    },
                    'erosions_ulcers': {
                        'score': min(3, max(0, uceis_score - (uceis_score // 3) - (uceis_score % 3))),
                        'confidence': float(confidence),
                        'description': ['无糜烂', '糜烂<5mm', '糜烂>5mm', '溃疡形成'][min(3, max(0, uceis_score - (uceis_score // 3) - (uceis_score % 3)))]
                    }
                }

            # 获取严重程度
            severity = SEVERITY_MAPPING.get(uceis_score, "未知")

            processing_time = time.time() - start_time

            logger.info(f"图像推理完成，UCEIS评分: {uceis_score}, 严重程度: {severity}, 置信度: {confidence:.2f}")
            logger.info(f"推理耗时: {processing_time:.2f}秒")

            return {
                'uceis_score': uceis_score,
                'severity': severity,
                'confidence': confidence,
                'detailed_scores': detailed_scores,
                'processing_time': processing_time
            }

    except Exception as e:
        logger.error(f"模型推理失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=f"模型推理失败: {str(e)}")

# 创建FastAPI应用
app = FastAPI(
    title="EndoSight-UC AI推理服务器",
    description="UCEIS评分专用AI模型推理服务",
    version="2.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    """应用启动时加载模型"""
    success = load_model()
    if not success:
        logger.error("模型加载失败，服务无法启动")
        raise RuntimeError("模型加载失败")

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "EndoSight-UC AI推理服务器运行中",
        "status": "ready",
        "model_loaded": model is not None,
        "device": str(device) if device else None
    }

@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "device": str(device) if device else None
    }

@app.get("/v1/health")
async def health_check_v1():
    """健康检查 v1"""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "device": str(device) if device else None,
        "model_type": "UCEISScoreModel",
        "uceis_score_range": "0-8",
        "supported_formats": ["jpg", "jpeg", "png", "bmp", "tiff"]
    }

@app.post("/v1/predict", response_model=PredictionResponse)
async def predict_endpoint(file: UploadFile = File(...)):
    """预测接口"""
    # 验证文件类型
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="请上传有效的图像文件")

    # 生成任务ID
    task_id = str(uuid.uuid4())

    try:
        # 保存临时文件
        temp_dir = Path("temp_uploads")
        temp_dir.mkdir(exist_ok=True)

        temp_file = temp_dir / f"{task_id}_{file.filename}"

        # 保存上传的文件
        with open(temp_file, "wb") as buffer:
            content = await file.read()
            buffer.write(content)

        logger.info(f"收到图像文件: {file.filename}")

        # 预处理图像
        image_tensor = preprocess_image(str(temp_file))

        # 模型推理
        result = predict_with_model(image_tensor)

        # 清理临时文件
        temp_file.unlink(missing_ok=True)

        return PredictionResponse(
            task_id=task_id,
            uceis_score=result['uceis_score'],
            severity=result['severity'],
            confidence=result['confidence'],
            detailed_scores=result['detailed_scores'],
            processing_time=result['processing_time']
        )

    except HTTPException:
        # 重新抛出HTTP异常
        raise
    except Exception as e:
        logger.error(f"处理请求时发生错误: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")

        # 清理临时文件
        if 'temp_file' in locals():
            temp_file.unlink(missing_ok=True)

        raise HTTPException(status_code=500, detail=f"服务器内部错误: {str(e)}")

@app.get("/v1/model/info")
async def model_info():
    """获取模型信息"""
    return {
        "model_type": "UCEISScoreModel",
        "model_config": model_config,
        "device": str(device) if device else None,
        "model_loaded": model is not None,
        "uceis_score_range": "0-8",
        "supported_formats": ["jpg", "jpeg", "png", "bmp", "tiff"]
    }

if __name__ == "__main__":
    logger.info("启动EndoSight-UC AI推理服务器...")
    uvicorn.run(
        "inference_server:app",
        host="0.0.0.0",
        port=5003,  # 使用5003端口避免冲突
        reload=False,
        log_level="info"
    )